
# https://leetcode.cn/problems/house-robber-iii/

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

# Method 1: dynamic programming solution ?
class Solution:
    def rob(self, root: TreeNode) -> int:
        #@cache
        def dfs(root):
            if root is None:
                return 0
            if root.left is None and root.right is None:
                return root.val
            a = dfs(root.left) + dfs(root.right)
            b = root.val
            if root.left:
                b += dfs(root.left.left) + dfs(root.left.right)
            if root.right:
                b += dfs(root.right.left) + dfs(root.right.right)
            return max(a, b)

        return dfs(root)

case1 = TreeNode(3, 
                 left = TreeNode(2, None, TreeNode(3)),
                 right = TreeNode(3, None, TreeNode(1)))
case2 = TreeNode(3,
                 left = TreeNode(4, TreeNode(1), TreeNode(3)),
                 right = TreeNode(5, None, TreeNode(1)))
s = Solution()
print(s.rob(case1))
print(s.rob(case2))